import json
import torch
from hta.trace_analysis import TraceAnalysis
import matplotlib.pyplot as plt

#tiaoshi json wenjian shifou baohan youxiao shijian
with open("/home/autox/anaconda3/python_learn/trace/trace.json") as f:
    trace_data = json.load(f)
    assert "traceEvents" in trace_data
    assert len(trace_data["traceEvents"]) > 0
    print(f"event num:{len(trace_data['traceEvents'])}")

trace_dir = "/home/autox/anaconda3/python_learn/trace"
analyzer = TraceAnalysis(trace_dir=trace_dir)
time_spent_df = analyzer.get_temporal_breakdown(visualize=True)
print(time_spent_df)
time_spent_df.plot(kind='bar')
plt.savefig('temporal_breakdown.png')

idle_time_df = analyzer.get_idle_time_breakdown()
print(idle_time_df)
# idle_time_df.plot(kind='bar')
# plt.savefig('idle_breakdown.png')

analyzer = TraceAnalysis(trace_dir=trace_dir)
analyzer.generate_trace_with_counters()